"""
核反应堆子通道分析程序 - 基于两流体模型

本程序实现了一个简化的核反应堆子通道分析工具，用于模拟和分析燃料棒束中冷却剂的流动特性。
采用两流体模型处理气液两相流动，考虑了轴向和横向的质量、动量传递。

主要功能：
1. 模拟子通道内的两相流动
2. 计算空泡率、压力分布
3. 分析液相和气相的速度场
4. 考虑子通道间的横向混合

物理模型：
- 两流体模型：分别处理气相和液相
- 考虑相间作用力
- 包含重力和压力梯度效应
- 简化的横向混合模型

数值方法：
- 采用显式时间推进格式
- 交错网格离散
- 迭代求解直至收敛

作者：[作者名称]
日期：[日期]
"""

import numpy as np
from CoolProp.CoolProp import PropsSI
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文显示
plt.rcParams['axes.unicode_minus'] = False  # 设置负号显示


class SubchannelAnalysis:
    """
    子通道分析的主类，实现了两流体模型的求解
    
    主要属性：
    - n_axial: 轴向节点数
    - n_channels: 子通道数量
    - channel_connections: 子通道连接关系
    - connection_coefficients: 子通道间的连接系数
    
    主要方法：
    - solve_mass_momentum_energy: 求解质量、动量和能量守恒方程
    - calculate_cross_flow: 计算横向流动
    - get_fluid_properties: 获取流体物性
    - calculate_heat_transfer_coefficient: 计算对流换热系数
    - calculate_phase_change_mass_flow: 计算相变质量流量
    - calculate_turbulent_mixing: 计算湍流混合和空泡漂移质量流量
    - calculate_wall_friction: 计算壁面摩擦阻力
    - calculate_interfacial_friction: 计算界面摩擦阻力
    """

    def __init__(self, n_axial=2000):
        """
        初始化子通道分析器
        
        参数：
        n_axial: int, 轴向节点数
        """
        # ====== 网格参数 ======
        self.n_axial = n_axial                    # 轴向节点数
        self.channel_length = 3.0                 # 通道总长度 (m)
        self.dz = self.channel_length / n_axial   # 轴向网格间距 (m)
        
        # ====== 几何参数 ======
        self.channel_diameter = 0.01              # 子通道等效直径 (m)
        self.channel_area = np.pi * (self.channel_diameter/2)**2  # 子通道横截面积 (m²)
        
        # ====== 子通道拓扑结构 ======
        # 定义子通道连接关系（根据图示）
        # 键：子通道编号，值：相邻子通道编号列表
        self.channel_connections = {
            1: [2, 3],      # 1号子通道与2、3号相邻
            2: [1, 3, 4],   # 2号子通道与1、3、4号相邻
            3: [1, 2, 4, 5],# 3号子通道与1、2、4、5号相邻
            4: [2, 3, 5, 6],# 4号子通道与2、3、5、6号相邻
            5: [3, 4, 6],   # 5号子通道与3、4、6号相邻
            6: [4, 5]       # 6号子通道与4、5号相邻
        }
        
        # 定义子通道间的连接系数（用于横向流动计算）
        # 键：(通道i, 通道j)，值：连接系数
        self.connection_coefficients = {
            (1, 2): 1, (2, 1): 1,  # 1-2通道间的连接
            (1, 3): 2, (3, 1): 2,  # 1-3通道间的连接
            (2, 3): 3, (3, 2): 3,  # 2-3通道间的连接
            (2, 4): 4, (4, 2): 4,  # 2-4通道间的连接
            (3, 4): 4, (4, 3): 4,  # 3-4通道间的连接
            (3, 5): 5, (5, 3): 5,  # 3-5通道间的连接
            (4, 5): 6, (5, 4): 6,  # 4-5通道间的连接
            (4, 6): 7, (6, 4): 7,  # 4-6通道间的连接
            (5, 6): 5, (6, 5): 5   # 5-6通道间的连接
        }
        
        # ====== 入口工况参数 ======
        self.inlet_pressure = 15.5e6  # 入口压力 (Pa)
        self.inlet_temp = 280         # 入口温度 (K)
        self.mass_flux = 1000         # 质量流量 (kg/m²s)
        
        # ====== 场变量初始化 ======
        self.n_channels = 6           # 子通道总数
        # 创建各个物理量的二维数组 (轴向节点 × 子通道数)
        self.void_fraction = np.zeros((n_axial, self.n_channels))    # 空泡率 (-)
        self.liquid_velocity = np.zeros((n_axial, self.n_channels))  # 液相速度 (m/s)
        self.vapor_velocity = np.zeros((n_axial, self.n_channels))   # 气相速度 (m/s)
        self.pressure = np.full((n_axial, self.n_channels), self.inlet_pressure)  # 压力场 (Pa)
        self.temperature = np.full((n_axial, self.n_channels), self.inlet_temp)   # 流体温度场 (K)
        self.temperature_wall = np.full((n_axial, self.n_channels), self.inlet_temp)  # 包壳温度场 (K)
        self.h_l = np.zeros((n_axial, self.n_channels))  # 液相焓 (J/kg)
        self.h_v = np.zeros((n_axial, self.n_channels))  # 气相焓 (J/kg)
        
        # 设置初始条件
        self._set_initial_conditions()
        
    def _set_initial_conditions(self):
        """
        设置计算初始条件
        
        包括：
        1. 入口边界条件的设置
        2. 沿程压力和温度的初始估计
        3. 流体物性的初始化
        """
        try:
            # 计算入口处的流体物性
            rho_l = PropsSI('D', 'P', self.inlet_pressure, 'T', self.inlet_temp, 'Water')  # 液相密度
            h_l = PropsSI('H', 'P', self.inlet_pressure, 'T', self.inlet_temp, 'Water')    # 液相焓
            h_v = PropsSI('H', 'P', self.inlet_pressure, 'Q', 1, 'Water')                  # 饱和蒸汽焓
            
            # 初始化每个子通道
            for i in range(self.n_channels):
                # 设置入口条件
                self.liquid_velocity[0, i] = self.mass_flux / rho_l  # 由质量流量计算入口速度
                self.vapor_velocity[0, i] = 0.0                      # 入口处无蒸汽
                self.void_fraction[0, i] = 0.0                       # 入口处无空泡
                self.h_l[0, i] = h_l                                 # 液相焓
                self.h_v[0, i] = h_v                                 # 气相焓
                
                # 沿程参数的初始估计
                for j in range(self.n_axial):
                    # 假设线性压降和温升
                    self.pressure[j, i] = self.inlet_pressure - j * 1e4  # 简单的压降估计
                    self.temperature[j, i] = self.inlet_temp + j * 0.1   # 简单的温升估计
                    self.temperature_wall[j, i] = self.inlet_temp + j * 0.2  # 包壳温度估计
                    
        except Exception as e:
            print(f"初始化失败: {str(e)}")
            raise
    
    def get_fluid_properties(self, p, t, phase):
        """
        获取指定状态点的流体物性
        
        参数：
        p: float, 压力 (Pa)
        t: float, 温度 (K)
        phase: str, 相态 ('liquid' 或 'vapor')
        
        返回：
        dict, 包含密度、粘度和焓值的字典
        """
        try:
            # 确保参数在CoolProp的有效范围内
            p = max(611.655, min(p, 2.2064e7))  # 水的三相点到临界点压力范围
            t = max(273.16, min(t, 647.096))    # 水的三相点到临界点温度范围
            
            if phase == 'liquid':
                # 计算液相物性
                return {
                    'density': PropsSI('D', 'P', p, 'T', t, 'Water'),    # 密度 (kg/m³)
                    'viscosity': PropsSI('V', 'P', p, 'T', t, 'Water'),  # 动力粘度 (Pa·s)
                    'enthalpy': PropsSI('H', 'P', p, 'T', t, 'Water')    # 比焓 (J/kg)
                }
            else:  # vapor
                # 计算气相物性
                return {
                    'density': PropsSI('D', 'P', p, 'Q', 1, 'Water'),    # 密度 (kg/m³)
                    'viscosity': PropsSI('V', 'P', p, 'Q', 1, 'Water'),  # 动力粘度 (Pa·s)
                    'enthalpy': PropsSI('H', 'P', p, 'Q', 1, 'Water')    # 比焓 (J/kg)
                }
        except Exception as e:
            print(f"警告：物性计算失败 (P={p}, T={t})")
            print(f"错误信息: {str(e)}")
            return None
            
    def calculate_cross_flow(self, i, j, k):
        """
        计算子通道间的横向流动
        
        参数：
        i: int, 轴向位置索引
        j: int, 当前子通道编号
        k: int, 相邻子通道编号
        
        返回：
        float, 横向质量流量 (kg/s)
        """
        # 检查是否存在连接
        if (j, k) not in self.connection_coefficients:
            return 0.0
        
        # 获取连接系数    
        s_jk = self.connection_coefficients[(j, k)]
        
        # 计算压力差驱动的横向流动
        dp = self.pressure[i, j-1] - self.pressure[i, k-1]  # 压力差
        cross_flow = s_jk * dp * 1e-3  # 横向流量与压力差成正比
        
        return cross_flow
    
    def calculate_heat_transfer_coefficient(self, alpha, T_wall, T_fluid, phase):
        """
        计算对流换热系数
        
        参数：
        alpha: float, 空泡率
        T_wall: float, 包壳温度 (K)
        T_fluid: float, 流体温度 (K)
        phase: str, 相态 ('liquid' 或 'vapor')
        
        返回：
        float, 对流换热系数 (W/m²K)
        """
        if phase == 'liquid':
            if T_wall < PropsSI('T', 'P', self.pressure, 'Q', 0, 'Water'):
                # 过冷沸腾
                F_gam = 0.1  # 假设值
                h_l = 1 - F_gam * np.exp((2 * self.pressure / 1260) * (T_wall - PropsSI('T', 'P', self.pressure, 'Q', 0, 'Water')) / (2722 * (T_wall - T_fluid)))
            else:
                # 饱和核态沸腾
                h_l = np.exp((2 * self.pressure / 1260) * (T_wall - PropsSI('T', 'P', self.pressure, 'Q', 0, 'Water')) / (2722 * (T_wall - T_fluid)))
            return h_l
        else:
            return 0  # 气相传热系数
    
    def calculate_phase_change_mass_flow(self, h_l, h_v, T_wall, T_fluid, phase):
        """
        计算相变质量流量
        
        参数：
        h_l: float, 液相焓 (J/kg)
        h_v: float, 气相焓 (J/kg)
        T_wall: float, 包壳温度 (K)
        T_fluid: float, 流体温度 (K)
        phase: str, 相态 ('liquid' 或 'vapor')
        
        返回：
        float, 相变质量流量 (kg/s)
        """
        if phase == 'liquid':
            C_l_sh = 100  # 假设值
            Gamma_l_sh = C_l_sh * (h_l - PropsSI('H', 'P', self.pressure, 'Q', 0, 'Water'))
            Gamma_l_sc = 0  # 假设无冷凝
            return Gamma_l_sh - Gamma_l_sc
        else:
            C_v_sh = 100  # 假设值
            Gamma_v_sh = C_v_sh * (PropsSI('H', 'P', self.pressure, 'Q', 1, 'Water') - h_v)
            Gamma_v_sc = 0  # 假设无冷凝
            return Gamma_v_sh - Gamma_v_sc
    
    def calculate_turbulent_mixing(self, alpha_i, alpha_j, equilibrium_alpha):
        """
        计算湍流混合和空泡漂移质量流量
        
        参数：
        alpha_i: float, 当前子通道空泡率
        alpha_j: float, 相邻子通道空泡率
        equilibrium_alpha: float, 平衡空泡率
        
        返回：
        float, 湍流混合质量流量 (kg/s)
        """
        C_e = 0.1  # 假设值
        L_gap = 0.001  # 假设值
        L_tur = 0.1  # 假设值
        rho_f = self.get_fluid_properties(self.pressure, self.temperature, 'liquid')['density']
        rho_g = self.get_fluid_properties(self.pressure, self.temperature, 'vapor')['density']
        return C_e * L_gap * L_tur * (rho_f - rho_g) * (alpha_j - alpha_i - (alpha_j - equilibrium_alpha))
    
    def calculate_wall_friction(self, u, phase):
        """
        计算壁面摩擦阻力
        
        参数：
        u: float, 流速 (m/s)
        phase: str, 相态 ('liquid' 或 'vapor')
        
        返回：
        float, 壁面摩擦力 (N/m²)
        """
        if phase == 'liquid':
            tau_wl = (self.pressure[i, j] - self.pressure[i-1, j]) / self.dz * (self.channel_area ** 2) * self.dz
            return tau_wl
        else:
            return 0  # 气相壁面摩擦阻力
    
    def calculate_interfacial_friction(self, v_l, v_v):
        """
        计算界面摩擦阻力
        
        参数：
        v_l: float, 液相速度 (m/s)
        v_v: float, 气相速度 (m/s)
        
        返回：
        float, 界面摩擦力 (N/m²)
        """
        C_lv = 0.1  # 假设值
        A_lv = self.channel_area  # 假设值
        return 0.5 * C_lv * self.get_fluid_properties(self.pressure, self.temperature, 'liquid')['density'] * (v_l - v_v) ** 2 * A_lv
    
    def solve_mass_momentum_energy(self):
        """
        求解质量、动量和能量守恒方程
        
        采用显式时间推进格式求解：
        1. 质量守恒方程用于更新空泡率和质量流量
        2. 动量方程用于更新速度场
        3. 能量方程用于更新温度场
        4. 考虑横向流动的影响
        """
        # 数值求解参数
        dt = 1e-4          # 时间步长 (s)
        max_iterations = 100000  # 最大迭代次数
        tolerance = 1e-6    # 收敛判据
        
        # 迭代求解
        for n in range(max_iterations):
            max_error = 0.0
            
            # 保存上一时间步的值用于计算误差
            old_void = np.copy(self.void_fraction)
            old_vl = np.copy(self.liquid_velocity)
            old_vv = np.copy(self.vapor_velocity)
            old_temp = np.copy(self.temperature)
            old_temp_wall = np.copy(self.temperature_wall)
            old_hl = np.copy(self.h_l)
            old_hv = np.copy(self.h_v)
            
            # 空间推进
            for i in range(1, self.n_axial):
                for j in range(self.n_channels):
                    # 获取当前位置的流体物性
                    props_l = self.get_fluid_properties(
                        float(self.pressure[i,j]), 
                        float(self.temperature[i,j]), 
                        'liquid'
                    )
                    props_v = self.get_fluid_properties(
                        float(self.pressure[i,j]), 
                        float(self.temperature[i,j]), 
                        'vapor'
                    )
                    
                    if props_l is None or props_v is None:
                        continue
                    
                    # 提取物性值    
                    rho_l = props_l['density']    # 液相密度
                    rho_v = props_v['density']    # 气相密度
                    h_l = props_l['enthalpy']     # 液相焓
                    h_v = props_v['enthalpy']     # 气相焓
                    
                    # 计算横向流动贡献
                    cross_flows = []
                    for k in self.channel_connections[j+1]:
                        w_jk = self.calculate_cross_flow(i, j+1, k)
                        cross_flows.append(w_jk)
                    
                    # 计算对流换热系数
                    h_coeff_l = self.calculate_heat_transfer_coefficient(self.void_fraction[i,j], self.temperature_wall[i,j], self.temperature[i,j], 'liquid')
                    h_coeff_v = self.calculate_heat_transfer_coefficient(self.void_fraction[i,j], self.temperature_wall[i,j], self.temperature[i,j], 'vapor')
                    
                    # 计算相变质量流量
                    Gamma_l = self.calculate_phase_change_mass_flow(h_l, h_v, self.temperature_wall[i,j], self.temperature[i,j], 'liquid')
                    Gamma_v = self.calculate_phase_change_mass_flow(h_l, h_v, self.temperature_wall[i,j], self.temperature[i,j], 'vapor')
                    
                    # 计算湍流混合质量流量
                    T_m = self.calculate_turbulent_mixing(self.void_fraction[i,j], self.void_fraction[i,j], self.void_fraction[i,j])
                    
                    # 计算壁面摩擦阻力
                    tau_wl = self.calculate_wall_friction(self.liquid_velocity[i,j], 'liquid')
                    tau_wv = self.calculate_wall_friction(self.vapor_velocity[i,j], 'vapor')
                    
                    # 计算界面摩擦阻力
                    tau_lv = self.calculate_interfacial_friction(self.liquid_velocity[i,j], self.vapor_velocity[i,j])
                    
                    # 求解质量守恒方程（显式格式）
                    new_alpha = self.void_fraction[i-1,j] + dt * (
                        -self.vapor_velocity[i-1,j] * (self.void_fraction[i,j] - self.void_fraction[i-1,j]) / self.dz  # 对流项
                        + sum(cross_flows) / (rho_v * self.channel_area)  # 横向流动项
                        + Gamma_l / (rho_l * self.channel_area)  # 相变质量流量
                        + T_m / (rho_l * self.channel_area)  # 湍流混合质量流量
                    )
                    new_alpha = np.clip(new_alpha, 0, 0.9999)  # 限制在物理范围内
                    
                    # 求解液相动量方程
                    new_vl = self.liquid_velocity[i-1,j] + dt * (
                        -self.liquid_velocity[i-1,j] * (self.liquid_velocity[i,j] - self.liquid_velocity[i-1,j]) / self.dz  # 对流项
                        - (1 / rho_l) * (self.pressure[i,j] - self.pressure[i-1,j]) / self.dz  # 压力梯度
                        - 9.81  # 重力项
                        - tau_wl / rho_l  # 壁面摩擦阻力
                        - tau_lv / rho_l  # 界面摩擦阻力
                    )
                    
                    # 求解气相动量方程
                    new_vv = self.vapor_velocity[i-1,j] + dt * (
                        -self.vapor_velocity[i-1,j] * (self.vapor_velocity[i,j] - self.vapor_velocity[i-1,j]) / self.dz  # 对流项
                        - (1 / rho_v) * (self.pressure[i,j] - self.pressure[i-1,j]) / self.dz  # 压力梯度
                        - 9.81  # 重力项
                        - tau_wv / rho_v  # 壁面摩擦阻力
                        - tau_lv / rho_v  # 界面摩擦阻力
                    )
                    
                    # 求解能量方程
                    new_hl = h_l + dt * (
                        -self.liquid_velocity[i,j] * (h_l - self.h_l[i-1,j]) / self.dz  # 对流项
                        + h_coeff_l * (self.temperature_wall[i,j] - self.temperature[i,j])  # 传热项
                        + Gamma_l * (h_v - h_l)  # 相变焓
                    )
                    new_hv = h_v + dt * (
                        -self.vapor_velocity[i,j] * (h_v - self.h_v[i-1,j]) / self.dz  # 对流项
                        + h_coeff_v * (self.temperature_wall[i,j] - self.temperature[i,j])  # 传热项
                        + Gamma_v * (h_l - h_v)  # 相变焓
                    )
                    
                    # 求解包壳温度场（一维圆柱热传导方程）
                    new_temp_wall = self.temperature_wall[i,j] + dt * (
                        (1 / (self.channel_diameter / 2)) * (self.channel_diameter / 2 * self.get_fluid_properties(self.pressure, self.temperature_wall[i,j], 'liquid')['viscosity'] * (self.temperature_wall[i,j] - self.temperature[i,j])) / self.dz  # 传导项
                        + Gamma_l * (h_v - h_l)  # 相变热
                        + Gamma_v * (h_l - h_v)  # 相变热
                    )
                    
                    # 更新场变量
                    self.void_fraction[i,j] = new_alpha
                    self.liquid_velocity[i,j] = new_vl
                    self.vapor_velocity[i,j] = new_vv
                    self.h_l[i,j] = new_hl
                    self.h_v[i,j] = new_hv
                    self.temperature_wall[i,j] = new_temp_wall
                    
                    # 计算流体温度
                    self.temperature[i,j] = (rho_l * h_l + rho_v * h_v) / (rho_l + rho_v)
                    
                    # 计算最大误差
                    error = max(
                        abs(new_alpha - old_void[i,j]),
                        abs(new_vl - old_vl[i,j]),
                        abs(new_vv - old_vv[i,j]),
                        abs(new_temp_wall - old_temp_wall[i,j]),
                        abs(new_hl - old_hl[i,j]),
                        abs(new_hv - old_hv[i,j])
                    )
                    max_error = max(max_error, error)
            
            # 检查收敛性
            if max_error < tolerance:
                print(f"计算在第{n}次迭代后收敛")
                break
                
            if n == max_iterations - 1:
                print("警告：计算未收敛")
    
    def run_simulation(self):
        """运行完整的模拟计算"""
        print("开始数值计算...")
        self.solve_mass_momentum_energy()
        print("计算完成")
    
    def plot_results(self):
        # 创建网格点
        x = np.arange(self.n_axial)
        y = np.arange(self.n_channels)
        X, Y = np.meshgrid(x, y)
        
        # 第一个图形窗口：6个子图
        fig = plt.figure(figsize=(20, 15))
        
        # 1. 空泡率分布
        ax1 = fig.add_subplot(231, projection='3d')
        surf1 = ax1.plot_surface(X, Y, self.void_fraction.T, cmap=cm.viridis)
        ax1.set_title('空泡率分布')
        ax1.set_xlabel('轴向位置')
        ax1.set_ylabel('子通道编号')
        ax1.set_zlabel('空泡率')
        fig.colorbar(surf1, ax=ax1, shrink=0.5, aspect=5)
        
        # 2. 液相速度分布
        ax2 = fig.add_subplot(232, projection='3d')
        surf2 = ax2.plot_surface(X, Y, self.liquid_velocity.T, cmap=cm.viridis)
        ax2.set_title('液相速度分布')
        ax2.set_xlabel('轴向位置')
        ax2.set_ylabel('子通道编号')
        ax2.set_zlabel('速度 (m/s)')
        fig.colorbar(surf2, ax=ax2, shrink=0.5, aspect=5)
        
        # 3. 气相速度分布
        ax3 = fig.add_subplot(233, projection='3d')
        surf3 = ax3.plot_surface(X, Y, self.vapor_velocity.T, cmap=cm.viridis)
        ax3.set_title('气相速度分布')
        ax3.set_xlabel('轴向位置')
        ax3.set_ylabel('子通道编号')
        ax3.set_zlabel('速度 (m/s)')
        fig.colorbar(surf3, ax=ax3, shrink=0.5, aspect=5)
        
        # 4. 压力分布
        ax4 = fig.add_subplot(234, projection='3d')
        normalized_pressure = (self.pressure - np.min(self.pressure)) / (np.max(self.pressure) - np.min(self.pressure))
        surf4 = ax4.plot_surface(X, Y, normalized_pressure.T, cmap=cm.viridis)
        ax4.set_title('压力分布（归一化）')
        ax4.set_xlabel('轴向位置')
        ax4.set_ylabel('子通道编号')
        ax4.set_zlabel('压力')
        fig.colorbar(surf4, ax=ax4, shrink=0.5, aspect=5)
        
        # 5. 流体温度分布
        ax5 = fig.add_subplot(235, projection='3d')
        surf5 = ax5.plot_surface(X, Y, self.temperature.T, cmap=cm.viridis)
        ax5.set_title('流体温度分布')
        ax5.set_xlabel('轴向位置')
        ax5.set_ylabel('子通道编号')
        ax5.set_zlabel('温度 (K)')
        fig.colorbar(surf5, ax=ax5, shrink=0.5, aspect=5)
        
        # 6. 包壳温度分布
        ax6 = fig.add_subplot(236, projection='3d')
        surf6 = ax6.plot_surface(X, Y, self.temperature_wall.T, cmap=cm.viridis)
        ax6.set_title('包壳温度分布')
        ax6.set_xlabel('轴向位置')
        ax6.set_ylabel('子通道编号')
        ax6.set_zlabel('温度 (K)')
        fig.colorbar(surf6, ax=ax6, shrink=0.5, aspect=5)
        
        # 添加视角和网格
        for ax in [ax1, ax2, ax3, ax4, ax5, ax6]:
            ax.view_init(elev=20, azim=45)  # 设置视角
            ax.grid(True)
        
        plt.tight_layout()
        plt.show()
        
        # 第二个图形窗口：焓值分布和沿轴向的变化曲线
        plt.figure(figsize=(15, 10))
        
        # 1. 液相焓分布
        ax7 = plt.subplot(221, projection='3d')
        surf7 = ax7.plot_surface(X, Y, self.h_l.T, cmap=cm.viridis)
        ax7.set_title('液相焓分布')
        ax7.set_xlabel('轴向位置')
        ax7.set_ylabel('子通道编号')
        ax7.set_zlabel('焓值 (J/kg)')
        plt.colorbar(surf7, ax=ax7, shrink=0.5, aspect=5)
        
        # 2. 气相焓分布
        ax8 = plt.subplot(222, projection='3d')
        surf8 = ax8.plot_surface(X, Y, self.h_v.T, cmap=cm.viridis)
        ax8.set_title('气相焓分布')
        ax8.set_xlabel('轴向位置')
        ax8.set_ylabel('子通道编号')
        ax8.set_zlabel('焓值 (J/kg)')
        plt.colorbar(surf8, ax=ax8, shrink=0.5, aspect=5)
        
        # 3. 沿轴向的压力变化
        plt.subplot(223)
        for i in range(self.n_channels):
            plt.plot(x, self.pressure[:, i], label=f'通道{i+1}')
        plt.title('沿轴向的压力变化')
        plt.xlabel('轴向位置')
        plt.ylabel('压力 (Pa)')
        plt.legend()
        plt.grid(True)
        
        # 4. 沿轴向的温度变化
        plt.subplot(224)
        for i in range(self.n_channels):
            plt.plot(x, self.temperature[:, i], label=f'通道{i+1}')
        plt.title('沿轴向的流体温度变化')
        plt.xlabel('轴向位置')
        plt.ylabel('温度 (K)')
        plt.legend()
        plt.grid(True)
        
        plt.tight_layout()
        plt.show()

if __name__ == "__main__":
    analysis = SubchannelAnalysis(n_axial=50)
    analysis.run_simulation()
    analysis.plot_results() 